﻿#include "FieldGraphGenerator.h"

#include <vtkTextProperty.h>
#include <vtkActor.h>
#include <vtkPointData.h>
#include <vtkCellData.h>
#include <vtkProperty.h>
#include <vtkScalarBarRepresentation.h>
#include <vtkScalarBarActor.h>
#include <vtkDoubleArray.h>
#include <vtkGlyph3D.h>
#include <vtkLookupTable.h>
#include <vtkPolyDataMapper.h>
#include <vtkDataSetMapper.h>
#include <vtkArrowSource.h>
#include <vtkScalarBarWidget.h>
#include <vtkMaskPoints.h>
#include <vtkLODActor.h>
#include <vtkDataObject.h>
#include <vtkDataSet.h>
#include <vtkPointSet.h>
#include <QDebug>
#include "PipelineManager.h"

namespace pst
{

    FieldGraphGenerator::FieldGraphGenerator(QObject* parent)
        : QObject{ parent }
        , m_inputData{ nullptr }
        , m_isScalarAttributeSet{ false }
        , m_geometryActor(vtkActor::New())
        , m_scalarFieldActor(vtkActor::New())
        , m_geometryMapper(vtkDataSetMapper::New())
        , m_scalarsFieldMapper(vtkDataSetMapper::New())
        , m_currentUsingScalarType(StructureType::POINTS)
        , m_currentAddedPointArrayIndex{ -1 }
        , m_currentAddedCellArrayIndex{ -1 }
        , m_currentScalarsArrayRange{ -1,0 }
        , m_scalarPipelineManager(new PipelineManager(this))
        , m_usingData(nullptr)
    {
        InitGeometryActor();
    }

    FieldGraphGenerator::~FieldGraphGenerator()
    {
        m_geometryActor->Delete();
        m_scalarFieldActor->Delete();

        m_geometryMapper->Delete();
        m_scalarsFieldMapper->Delete();
    }

    void FieldGraphGenerator::SetInputDataSet(vtkDataSet* data)
    {
        m_inputData = data;
        if (m_usingData == nullptr) //确保m_usingData指针地址不变
        {
            m_usingData = vtkSmartPointer<vtkDataSet>::Take(m_inputData->NewInstance());
        }
        //PrintInputDataSetInformation();
    }

    void FieldGraphGenerator::PrintInputDataSetInformation()
    {
        qDebug() << "\n\n=============================";
        qDebug() << "Input data information:\n";

        qDebug() << "points num = " << m_inputData->GetNumberOfPoints();
        qDebug() << "Points data array information:";
        for (int i = 0; i < m_inputData->GetPointData()->GetNumberOfArrays(); ++i)
        {
            auto array = m_inputData->GetPointData()->GetArray(i);
            PrintArrayInformation(array);
        }

        qDebug() << "cell num = " << m_inputData->GetNumberOfCells();
        qDebug() << "Cells data array information:";
        for (int i = 0; i < m_inputData->GetCellData()->GetNumberOfArrays(); ++i)
        {
            auto array = m_inputData->GetCellData()->GetArray(i);
            PrintArrayInformation(array);
        }
        qDebug() << "=============================\n\n";
    }

    vtkDataSet* FieldGraphGenerator::GetInputData()
    {
        return m_inputData;
    }

    QList<FieldGraphGenerator::ArrayInformation>
        FieldGraphGenerator::GetPointsArrayIndexAndNames()
    {
        QList<FieldGraphGenerator::ArrayInformation>  arrayInfoVector;
        FieldGraphGenerator::ArrayInformation tempInfo;

        for (int i = 0; i < m_inputData->GetPointData()->GetNumberOfArrays(); ++i)
        {
            tempInfo.arrayIndex = i;
            tempInfo.arrayName = m_inputData->GetPointData()->GetArrayName(i);
            tempInfo.arrayComponent = m_inputData->GetPointData()->GetArray(i)->GetNumberOfComponents();
            tempInfo.type = StructureType::POINTS;
            arrayInfoVector.push_back(tempInfo);
        }
        return arrayInfoVector;
    }

    QList<FieldGraphGenerator::ArrayInformation>
        FieldGraphGenerator::GetCellsArrayIndexAndNames()
    {
        QList<FieldGraphGenerator::ArrayInformation>  arrayInfoVector;
        FieldGraphGenerator::ArrayInformation tempInfo;

        for (int i = 0; i < m_inputData->GetCellData()->GetNumberOfArrays(); ++i)
        {
            tempInfo.arrayIndex = i;
            tempInfo.arrayName = m_inputData->GetCellData()->GetArrayName(i);
            tempInfo.arrayComponent = m_inputData->GetCellData()->GetArray(i)->GetNumberOfComponents();
            tempInfo.type = StructureType::CELLS;
            arrayInfoVector.push_back(tempInfo);
        }
        return arrayInfoVector;
    }

    bool FieldGraphGenerator::SetActiveScalarAttribute(int index, 
        const StructureType type, const int compIndex, bool isLog, int logFactor, bool isNormalized)
    {
        
        m_errorMsg.clear(); 
        m_usingData->DeepCopy(m_inputData);

        //移除上次额外添加的数组
        //m_usingData->GetPointData()->RemoveArray(m_currentAddedPointArrayIndex);
        //m_usingData->GetCellData()->RemoveArray(m_currentAddedCellArrayIndex);

        vtkDataArray* _activeArray = nullptr;
        switch (type)
        {
        case StructureType::POINTS:
        {
            if (index >= m_usingData->GetPointData()->GetNumberOfArrays() || index < -1)
            {
                qDebug() << "Error! Index illegal!";
                m_errorMsg = "Error! Index illegal!";
                return false;
            }

            auto activeScalarArray = m_usingData->GetPointData()->GetArray(index);
            const auto nComponents = activeScalarArray->GetNumberOfComponents();
            _activeArray = activeScalarArray;

            //单组分array,忽略组分提取，直接设置
            if (nComponents == 1)
            {
                m_usingData->GetPointData()->SetActiveAttribute(
                    index, vtkDataSetAttributes::AttributeTypes::SCALARS);

            }
            else
            {
                if (compIndex < -1 || compIndex >= nComponents)
                {
                    qDebug() << "Error! Extracting component index is " << compIndex
                        << ", current array component is " << nComponents;
                    m_errorMsg = "Error! Extracting component index is " + QString::number(compIndex)
                        + ", current array component is " + QString::number(nComponents);
                    return false;
                }

                vtkNew<vtkDoubleArray> magArray;
                bool isSuccess = GetSingleComponent(activeScalarArray, compIndex, magArray);
                if (!isSuccess)
                {
                    qDebug() << "Extract component failed!";
                    m_errorMsg = "Extract component failed!";
                    return false;
                }

                m_currentAddedPointArrayIndex = m_usingData->GetPointData()->AddArray(magArray);
                m_usingData->GetPointData()->SetActiveAttribute(
                    m_currentAddedPointArrayIndex, vtkDataSetAttributes::AttributeTypes::SCALARS);
                _activeArray = magArray;
            }
            m_scalarsFieldMapper->SetScalarModeToUsePointData();
            m_currentUsingScalarType = StructureType::POINTS;

            m_usingData->GetCellData()->SetActiveScalars(nullptr);//取消激活cellData
            break;
        }

        case StructureType::CELLS:
        {
            if (index >= m_usingData->GetCellData()->GetNumberOfArrays() || index < 0)
            {
                qDebug() << "Error! Index illegal!";
                m_errorMsg = "Error! Index illegal!";
                return false;
            }


            vtkDataArray* array = m_usingData->GetCellData()->GetArray(index);
            //vtkDataArray* activeScalarArray = vtkDoubleArray::SafeDownCast(array);
            const auto nComponents = array->GetNumberOfComponents();
            _activeArray = array;

            //单组分array,忽略组分提取，直接设置
            if (nComponents == 1)
            {
                m_usingData->GetCellData()->SetActiveAttribute(
                    index, vtkDataSetAttributes::AttributeTypes::SCALARS);
            }
            else
            {
                if (compIndex < -1 || compIndex >= nComponents)
                {
                    qDebug() << "Error! Extracting component index is " << compIndex
                        << ", current array component is " << nComponents;
                    m_errorMsg = "Error! Extracting component index is " + QString::number(compIndex)
                        + ", current array component is " + QString::number(nComponents);
                    return false;
                }

                vtkNew<vtkDoubleArray> magArray;
                bool isSuccess = GetSingleComponent(array, compIndex, magArray);
                if (!isSuccess)
                {
                    qDebug() << "Extract component failed!";
                    m_errorMsg = "Extract component failed!";
                    return false;
                }

                _activeArray = magArray;
                m_currentAddedCellArrayIndex = m_usingData->GetCellData()->AddArray(magArray);
                m_usingData->GetCellData()->SetActiveAttribute(
                    m_currentAddedCellArrayIndex, vtkDataSetAttributes::AttributeTypes::SCALARS);

            }
            m_scalarsFieldMapper->SetScalarModeToUseCellData();
            m_currentUsingScalarType = StructureType::CELLS;

            m_usingData->GetPointData()->SetActiveScalars(nullptr);//取消激活pointData
            break;
        }

        default:
            qDebug() << "Wrong structure type";
            m_errorMsg = "Wrong structure type";
            return false;
            break;
        }

        // 取对数和归一化
        if (_activeArray)
        {
            if (isNormalized)
            {
                normalizeTheArray(vtkDataArray::SafeDownCast(_activeArray));
            }
            if (isLog)
            {
               auto isSuc =  logTheArray(vtkDataArray::SafeDownCast(_activeArray), logFactor);
               if (!isSuc)
               {
                   return false;
               }
            }
        }
        m_currentScalarsArrayRange[0] = _activeArray->GetRange()[0];
        m_currentScalarsArrayRange[1] = _activeArray->GetRange()[1];

        m_isScalarAttributeSet = true;
        return true;
    }

    PipelineManager* FieldGraphGenerator::GetScalarPipelineManager()
    {
        return m_scalarPipelineManager;
    }

    std::array<double, 2> FieldGraphGenerator::GetActiveFieldOriginalRange()
    {
        return m_currentScalarsArrayRange;
    }

    vtkActor* FieldGraphGenerator::GetScalarsFieldActor()
    {
        return m_scalarFieldActor;
    }

    void FieldGraphGenerator::SetGeometryActorColor(
        const double r, const double g, const double b)
    {
        m_geometryActor->GetProperty()->SetColor(r, g, b);
    }

    void FieldGraphGenerator::UpdateScalarFieldActor()
    {
        if (m_isScalarAttributeSet)
        {
            vtkDataObject* processedData = m_scalarPipelineManager->Process(m_usingData);
            vtkDataSet* dataSet = vtkDataSet::SafeDownCast(processedData);
            if (dataSet)
            {
                m_scalarsFieldMapper->SetInputData(dataSet);
            }
            else
            {
                qDebug() << "Error: processedData cannot be cast to vtkDataSet.";
                return;
            }
            //m_scalarsFieldMapper->SetLookupTable(m_scalarsFieldLookupTable);
            m_scalarsFieldMapper->SetInterpolateScalarsBeforeMapping(true);
            m_scalarsFieldMapper->SetColorModeToMapScalars();
            m_scalarsFieldMapper->ScalarVisibilityOn();
            m_scalarsFieldMapper->SetUseLookupTableScalarRange(1);
            m_scalarsFieldMapper->Update();
            m_scalarFieldActor->SetMapper(m_scalarsFieldMapper);
            m_scalarFieldActor->GetProperty()->SetRepresentationToSurface();
            //        m_scalarFieldActor->GetProperty()->SetRepresentationToWireframe();

            //m_scalarsScalarBarWidget->On();
        }
        //else
        //{
        //    qDebug() << "A scalar field should be activated first! ";
        //}
    }

    void FieldGraphGenerator::UpdateGeometryActor()
    {
        vtkDataObject* processedData = m_scalarPipelineManager->Process(m_inputData);
        vtkDataSet* dataSet = vtkDataSet::SafeDownCast(processedData);
        if (dataSet)
        {
            m_geometryMapper->SetInputData(dataSet);
        }
        m_geometryMapper->Update();
    }

    vtkActor* FieldGraphGenerator::GetGeometryActor()
    {
        return m_geometryActor;
    }

    QString FieldGraphGenerator::getErrorMessage() const
    {
        return m_errorMsg;
    }

    vtkSmartPointer<vtkDataSet> FieldGraphGenerator::getUsingData() const
    {
        return m_usingData;
    }

    void FieldGraphGenerator::InitGeometryActor()
    {
        m_geometryMapper->ScalarVisibilityOff();
        m_geometryActor->SetMapper(m_geometryMapper);
        m_geometryActor->GetProperty()->SetColor(
            0.972549019, 0.9725490196, 0.64313725490);
    }

    void FieldGraphGenerator::PrintArrayInformation(vtkDataArray* array)
    {
        qDebug() << "name = " << array->GetName()
            << "\nsize = " << array->GetNumberOfTuples()
            << "\ncomponent = " << array->GetNumberOfComponents()
            << "\nscalar range = " << array->GetRange()[0]
            << " " << array->GetRange()[1] << "\n\n";
    }

    bool FieldGraphGenerator::GetSingleComponent(vtkDataArray* array,
        const int componentIndex, vtkDataArray* returnArray)
    {
        auto nComponents = array->GetNumberOfComponents();
        qDebug() << " nComponents = " << nComponents;

        //检查下标合法性，-1代表取magnitude
        if (componentIndex < -1 || componentIndex >= nComponents)
        {
            return false;
        }

        //取单个组分的属性的0组分，即返回自己
        if (componentIndex == 0 && nComponents == 1)
        {
            returnArray->DeepCopy(vtkDataArray::SafeDownCast(array));
            return true;
        }

        const auto nTuples = array->GetNumberOfTuples();
        qDebug() << " nTuples = " << nTuples;

        vtkNew<vtkDoubleArray> singleArray;
        singleArray->SetNumberOfComponents(1);
        singleArray->SetNumberOfTuples(nTuples);

        //取magnitude
        if (componentIndex == -1)
        {
            qDebug() << "magnitude ";
            singleArray->SetName("Magnitude");
            for (vtkIdType i = 0; i < nTuples; ++i)
            {
                double mag = 0;
                for (int j = 0; j < nComponents; ++j)
                {
                    double tmp = array->GetComponent(i, j);
                    mag += tmp * tmp;
                }
                mag = sqrt(mag);
                singleArray->InsertTuple1(i, mag);
            }
        }
        else //取现有的单个组件
        {
            qDebug() << "extract comp =  " << componentIndex;

            if (componentIndex == 0)
            {
                singleArray->SetName("X");
            }
            else if (componentIndex == 1)
            {
                singleArray->SetName("Y");
            }
            else if (componentIndex == 2)
            {
                singleArray->SetName("Z");
            }
            else
            {
                singleArray->SetName(std::to_string(componentIndex).c_str());
            }


            for (vtkIdType j = 0; j < nTuples; ++j)
            {
                singleArray->InsertTuple1(j, array->GetComponent(j, componentIndex));
                //if(j<50)
                //{
                //   qDebug() <<"bb= "<<array->GetComponent(j,componentIndex)<<std::endl;
                //}
            }
        }
        returnArray->DeepCopy(singleArray);
        return true;
    }

    bool FieldGraphGenerator::logTheArray(vtkDataArray* array, int logFactory)
    {
        if (array == nullptr)
        {
            m_errorMsg = "The array is nullptr, please check the data and try again";
            return false;
        }
        auto rowNum = array->GetNumberOfTuples();
        for (int i = 0; i < rowNum; ++i)
        {
            auto value = array->GetVariantValue(i).ToDouble();
            if (qFuzzyCompare(value, 0.0) || (!qFuzzyIsNull(value) && value < 0.0))
            {
                m_errorMsg = "数值存在0或负数，请检查数据！";
                return false;
            }
            array->SetVariantValue(i, logFactory * std::log10(value));
        }
        array->Modified();
        return true;
    }

    void FieldGraphGenerator::normalizeTheArray(vtkDataArray* array)
    {
        if (array == nullptr)
        {
            return;
        }

        //auto maxValue = array->GetRange()[1];

        vtkIdType nTuples = array->GetNumberOfTuples();
        int nComp = array->GetNumberOfComponents();
        double maxValue = 0.0;//改为绝对值的最大值

        for (vtkIdType i = 0; i < nTuples; ++i)
        {
            for (int c = 0; c < nComp; ++c)
            {
                double v = array->GetComponent(i, c);
                double av = std::abs(v);
                if (av > maxValue)
                {
                    maxValue = av;
                }
            }
        }

        auto rowNum = array->GetNumberOfTuples();
        for (int i = 0; i < rowNum; ++i)
        {
            auto value = array->GetVariantValue(i).ToDouble();
            array->SetVariantValue(i, value / maxValue);
        }
        array->Modified();
    }

}//namespace pst
